import os, sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from .base import BaseDebateAgent, DebateState
from my_tools.debate_evidence_retriever import DebateEvidenceTool
from langchain_core.prompts import ChatPromptTemplate
from langgraph.graph import END, Graph, START


class RAGDebateAgent(BaseDebateAgent):
    """基于LangGraph的实现（含证据工具）"""

    def __init__(self, llm, evidence_tool: DebateEvidenceTool = None):
        super().__init__(llm)
        self.role_prompt = ""
        self.evidence_tool = evidence_tool or DebateEvidenceTool()
        self.evidence_text = ""
        self._build_workflow()

    def _build_workflow(self):
        """构建LangGraph工作流"""
        self.workflow = Graph()

        # 定义节点
        self.workflow.add_node("decide_evidence", self._decide_evidence)
        self.workflow.add_node("retrieve_evidence", self._retrieve_evidence)
        self.workflow.add_node("generate_speech", self._generate_speech)

        # 设置条件分支
        self.workflow.add_conditional_edges(
            START,
            self._decide_evidence,
            {"evidence_needed": "retrieve_evidence", "no_evidence": "generate_speech"}
        )

        # 设置常规边
        # self.workflow.add_edge(START, "decide_evidence")
        self.workflow.add_edge("retrieve_evidence", "generate_speech")
        self.workflow.add_edge("generate_speech", END)

        # 编译工作流
        self.compiled_workflow = self.workflow.compile()

    def _decide_evidence(self, state: DebateState):
        """决策节点逻辑"""
        prompt = ChatPromptTemplate.from_template(
            self.role_prompt + """
            现在你需要判断是否需要使用证据或数据来支撑你的观点。如果需要使用证据或数据，请根据向量数据库的特性，直接回答查询关键词(关键词不要超过10个字)，格式如：‘查询：……’，不需要其它说明；如果不需要，请回答“不需要”。
        """)
        chain = prompt | self.llm
        decision = chain.invoke({}).content
        state["llm_thought"] = self._remove_think_tags(decision)
        print(state["llm_thought"])
        return "no_evidence" if "不需要" in decision else "evidence_needed"

    def _retrieve_evidence(self, state: DebateState):
        """证据检索逻辑"""
        evidence = self.evidence_tool.tool.run({"claim": self.extract_query_content(state["llm_thought"])})
        self.evidence_text = evidence
        print(self.evidence_text)

    def _generate_speech(self, state: DebateState):
        """辩词生成逻辑"""
        temp_prompt = """
        请根据历史发言摘要,当前获取到的证据和你当前的角色、任务，生成你的发言。
        你的发言要求：
            1. 使用口语化表述
            2. 一定确保不解释、描述或总结你的思考过程和发言意图。
            3. 每个论点必须包含具体理论、事实、案例或数据
            4. 不要提及数据来源及类似需调整内容的说明性文字
            4. 科学严谨、简洁清晰
        """

        if self.evidence_text == "":
            prompt = ChatPromptTemplate.from_template(self.role_prompt + temp_prompt)
            chain = prompt | self.llm
            response = chain.invoke({}).content
        else:
            prompt = ChatPromptTemplate.from_template(
            self.role_prompt + """
            以下证据是你需要使用的证据或数据:
            {evidence_text}
        请根据历史发言摘要,当前获取到的证据和你当前的角色、任务，生成你的发言。
        你的发言要求：
            1. 使用口语化表述
            2. 一定确保不解释、描述或总结你的思考过程和发言意图。
            3. 每个论点必须包含具体理论、事实、案例或数据
            4. 不要提及数据来源及类似需调整内容的说明性文字
            5. 科学严谨、简洁清晰
            6. 不要重复我方已经使用的证据。
        """
                                                  )
            chain = prompt | self.llm
            response = chain.invoke({"evidence_text": self.evidence_text}).content
        self.evidence_text = ""
        return {"speech": response}

    def generate_role_prompt(self, state: DebateState):
        self.role_prompt = f"""
        现正在进行辩题为：{state['topic']} 的辩论赛。
        你是辩论赛{state['side']}的成员，你的角色是{state['role']}。
        你的角色提示：{state['role_prompt']}。
        当前辩论阶段：{state['stage']}
        你的当前任务：{state['task']}
        辩论历史摘要：
        {state['history']}
        """

    def extract_query_content(self, text: str) -> str:
        """提取字符串中『查询：』之后的内容"""
        query_start = text.find("查询：")
        if query_start != -1:
            # 找到后截取内容并去除前后空格
            return text[query_start + 3:].strip()
        return text

    def _remove_think_tags(self, text):
        start_index = text.find("<think>")
        end_index = text.find("</think>")

        if start_index != -1 and end_index != -1:
            end_index += len("</think>")
            r_text = text[:start_index] + text[end_index:]
            return r_text.strip()
        return text

    def invoke(self, side: str, topic: str, stage: str, role: str,
               role_prompt: str, task: str, history: list = None) -> str:
        state: DebateState = {
            "side": side,
            "topic": topic,
            "stage": stage,
            "role": role,
            "role_prompt": role_prompt,
            "task": task,
            "history": history or [],
            "llm_thought": "",
            "speech": ""
        }
        self.generate_role_prompt(state)
        result = self.compiled_workflow.invoke(state)
        return self._remove_think_tags(result["speech"])


if __name__ == "__main__":
    from customize.get_ollama import GetOllama

    llm = GetOllama(ip="10.12.8.21:11434", model_name="deepseek-r1:14b")()
    llm2 = GetOllama(model_name="qwen2.5:3b")()
    debate_agent = RAGDebateAgent(llm=llm, evidence_tool=DebateEvidenceTool(llm=llm))
    print(debate_agent.invoke(side="正方", topic="手机应该进入中学校园", stage="论证阶段", role="正方辩手",
                              role_prompt="", task="论证观点"))
